#! /usr/bin/python

'''This script reads in a set of spike time files exported from the sortpy program and makes a multi-raster plot. In addition an option for a histogram is included'''

import os
import scipy.io
from scipy import array
import sys
import numpy
import pylab

import os

files = os.listdir('.')
datafiles = []
countlist = []
count = 1
filenames = []
all_spikes = []

def extract_data(filename):
    return numpy.loadtxt(filename)

pylab.figure()
total_time = 0
for file in sorted(files):
    if file.endswith('detection.txt'):
        spike_xs = extract_data(file)
        all_spikes.extend(spike_xs)
        spike_ys = numpy.ones(len(spike_xs))*(count+1)
        pylab.subplot(211)
        pylab.plot(spike_xs, spike_ys, linewidth=0, marker='|',
                    markersize=8, markeredgewidth=.5, color='k')
        
        count += 1
        filenames.append(file.split('-')[0])
        countlist.append(count)
    if file.endswith('raw_traces.txt'):
        if total_time == 0:
            a = extract_data(file)
            total_time = a[0][-1]
pylab.ylim(1,count+1)
pylab.xlim(0,total_time)
ax = pylab.gca()
pylab.xlabel('time(ms)')
pylab.ylabel('Trial #', fontsize=15)        
ax.set_yticks(countlist)
ax.set_yticklabels(filenames)
#pylab.suptitle('Bar Orientation Tuning',fontsize = 22)
pylab.show()
    
##############################################################################
################################ HISTOGRAM ###################################
##############################################################################

def bin_spikes(spike_list, total_time, bin_width=100.0,bin_alignment="middle"):
    high_end = bin_width

    low  = 0.0
    high = 0.0
    def is_between(value):
       return low <= value <= high
    
    bins = numpy.arange(0.0, total_time, bin_width)
    spikes = []
    for i in xrange(len(bins)-1):
        low = bins[i]
        high = bins[i+1]
        spike_count = len(filter(is_between, spike_list))
        spikes.append(spike_count)
    # Take care of final bin
    low = bins[-1]
    high = total_time
    spike_count = len(filter(is_between, spike_list))
    spikes.append(spike_count)
    return spikes, bins
#axes = pylab.subplot(212)
number_trials = count - 1
bin_width = 100.0
spike_norm_factor = 1000.0/(bin_width*number_trials) 
spikes, bins = bin_spikes(all_spikes, total_time, bin_width=bin_width)
for bin, spike_count in zip(bins,spikes):
    pylab.bar(bin, spike_count*spike_norm_factor, width=bin_width)
pylab.xlabel('time (ms)', fontsize=15)
pylab.ylabel('Hz/trial', fontsize=15)
pylab.show()

